#ifndef __Analysis_h__
#define __Analysis_h__
#include <iostream>
#include <iomanip>
#include <sstream>
#include <cstdlib>
#include <map>
#include <algorithm>
#include <fstream>
#include <time.h>
#include "CmdLine.hh"
#include "TFile.h"
#include "TChain.h"

#include "TH1F.h"
#include "TH2F.h"

#include "fastjet/ClusterSequence.hh"
#include "fastjet/tools/Pruner.hh"
#include "fastjet/tools/Filter.hh"

using namespace std;
using namespace fastjet;

//a wrapper over histograms to pretend seg faults

class Plot
{
 public:

  Plot(TH1* a=0) : p(a){}
  
  TH1* operator->()
  {
    if(p==NULL)
      {
	cout<<"Trying to access null histogram! recheck your code"<<endl;
	throw;
      }
    else 
      return p;
  }
 
  Plot& operator=(TH1* h)
  { p=h; return *this; }

 public:
  TH1* p;

};

class Plot2
{
 public:

  Plot2(TH2* a=0) : p(a){}
  
  TH2* operator->()
  {
    if(p==NULL)
      {
	cout<<"Trying to access null histogram! recheck your code"<<endl;
	throw;
      }
    else 
      return p;
  }
 
  Plot2& operator=(TH2* h)
  { p=h; return *this; }

 private:
  TH2* p;
};

class Analysis
{
 public:
  //histograms to be processed
  map<string, Plot> hist;
  map<string, Plot2> hist2;
  TChain tree;
  PseudoJet met;
  const static float jet_minpt=5.0;



  //jets
  vector<PseudoJet> InputJets, jets;
  vector<PseudoJet> leps;

  //quick variables
  vector<Double_t> ttop_pt;
  ClusterSequence* clustSeq;
  


  ~Analysis(void){}
  //user functions
  void hist_init(){}
  void hist_init2(){}
  void analyze(){}
  //determine whether event
  //passes our selections
  bool cut(){return true;}

  //get the index for cell

  static int phi_i(double phi)
    {
      if(phi<0) phi+=6.28318531;
      return static_cast<int>(phi/stepsize_);
    }
    
  static int eta_i(double eta)
    {return static_cast<int>((eta-min_eta)/stepsize_);}

 
  void makejets(int option=0, JetAlgorithm opt_j=antikt_algorithm, 
		double opt_r=1.0, double minpt=50.0)
  {
    //reset all the jets
    reset();

    using namespace fastjet;
    //declare fastjet algorithm
    //antikt_algorithm, cambridge_algorithm or kt_algorithm
    JetAlgorithm jet_algorithm=antikt_algorithm;
    
    if(s_algorithm=="kt")
      jet_algorithm=kt_algorithm;
    else if(s_algorithm=="cambridge")
      jet_algorithm=cambridge_algorithm;
    
    
    JetDefinition jetDef(jet_algorithm, r);
    if(option!=0)
      {
	jetDef=JetDefinition(opt_j, opt_r);
	//cout<<"radius is : "<<opt_r<<endl;
      }

    for(Int_t i=0; i<size; i++)
    if(Pt[i]>1e-20)
      {
	if(abs(PID[i])==6)
	  ttop_pt.push_back(sqrt(Px[i]*Px[i]+Py[i]*Py[i]));


	if(Status[i]!=1)
	  continue;

	int c_pid=abs(PID[i]);

	if(c_pid==12 || c_pid==14 || c_pid==16)
	  continue;

	if(c_pid==1000022 || c_pid==2000012 ||
	   c_pid==2000014 || c_pid==2000016 ||
	   c_pid==1000039 || c_pid==5000039 ||
	   c_pid==4000012 || c_pid==9900012 || 
	   c_pid==9900014 || c_pid==9900016 || c_pid==39)
	  continue;


	PseudoJet j(Px[i], Py[i], 
		    Pz[i], E[i]);

	if(j.m() > 20.0) continue;
	
	j.set_cached_rap_phi(Eta[i], Phi[i]);

	if(fabs(Eta[i])>3)
	  continue;


	if(j.pt() > 100.0)
	if(abs(PID[i])==11 || 
	   abs(PID[i])==13 )
	  leps.push_back(j);
	
	  if(fabs(Eta[i])>2.5)
	    if(PID[i]==22 ||
	       abs(PID[i])==11 || 
	       abs(PID[i])==13)
	      continue;
	  

	  if(b_pixel)
	    cell_E[eta_i(Eta[i])][phi_i(Phi[i])] += E[i];

	  else
	    InputJets.push_back(j);
	  /*
	  cout<<"par: "<<PID[i]<<", status: "<<Status[i]<<endl;
	  cout<<j.pt()<<","<<j.m()<<","<<j.eta()<<","<<j.phi()<<endl;
	  */
	  if(!b_pixel)
	    met-=j;
      }
      

    //do pixelation
    if(b_pixel)
      {
	//cout<<"pixelllllledddd!!!"<<endl;
	for(int i=0; i<cell_eta_size; i++)
	for(int j=0; j<cell_phi_size; j++)
	  {
	    if(cell_E[i][j] > 0)
	      {
		InputJets.push_back(cell_vec[i][j]);
		InputJets.back() *= cell_E[i][j];
		met-=InputJets.back();
	      }
	  }
	for(int i=0; i<leps.size(); i++)
	  {
	    met-= (cell_vec[ eta_i(leps[i].eta()) ][ phi_i(leps[i].phi()) ]
		   *leps[i].e());
	  }
      }


    //minpt = 5
    clustSeq=new ClusterSequence(InputJets, jetDef);
    jets = sorted_by_pt(clustSeq->inclusive_jets(minpt));

    //apply isolation to jets
    vector<PseudoJet>::iterator iter = jets.begin();
    while (iter != jets.end())
      {
	bool b_iso=false;
	double c_pt=iter->pt();
	
	//lepton isolation
	for(int k=0; k<leps.size(); k++)
	  if(2.0* leps[k].pt() > c_pt)
	    if(leps[k].delta_R(*iter) < 0.2)
	      {b_iso=true; break;}
	
	//remove jet if close to lepton
	if(b_iso) 
	  jets.erase(iter);
	else
	  ++iter;
	//cout<<"ERASED JETS"<<endl;
      }
      

    hist["met"]->Fill(met.pt(),weight);

    //if pruning is checked, do pruning
    if(b_prune)
      {
	//cout<<"WARNING: pruning"<<endl;
	double zcut=1.5;
	double Rcut_factor=1.0;
	Pruner pruner(kt_algorithm, zcut, Rcut_factor);
	
	for(int i=0; i<4 && i<jets.size(); i++)
	  {
	    jets[i]=pruner(jets[i]);
	  }
      }

  }


  void makejets(JetAlgorithm opt_j, double r, double minpt=50.0)
  {
    makejets(1, opt_j, r, minpt);
  }
  
 Analysis(const CmdLine& cmdline): tree("STDHEP"), cmdline(cmdline)
    {
      input_filename = cmdline.value<string>("-i");
      output_filename = cmdline.value<string>("-o");
      s_algorithm =cmdline.value<string>("-a", "antikt");
      
      if(s_algorithm=="kt")
	{cout<<"OPTION: kt algorithm"<<endl;}
      else if(s_algorithm=="cambridge")
	{cout<<"OPTION: cambridge algorithm"<<endl;}
      else
	{cout<<"OPTION: default antikt algorithm"<<endl;}

      //initialize cells for pixelization
      for(int i=0; i<cell_eta_size; i++)
      for(int j=0; j<cell_phi_size; j++)
	{
	  double cell_eta=min_eta+(i+0.5)*stepsize_;
	  double cell_phi=(j+0.5)*stepsize_;
	  
	  double cell_pt=1.0/cosh(cell_eta);
	  double cell_pz=tanh(cell_eta);
	  double cell_px=cell_pt*cos(cell_phi);
	  double cell_py=cell_pt*sin(cell_phi);
	  
	  cell_vec[i][j].
	    reset_momentum(cell_px, cell_py, cell_pz, 1.0);
	  cell_vec[i][j].
	    set_cached_rap_phi(cell_eta, cell_phi);
	}

      mass = cmdline.value<int>("-m",500);
      r = cmdline.value<double>("-r",1.2);
      b_trim = ! cmdline.present("-notrim");
      b_pixel = ! cmdline.present("-nopixel");
      b_prune = cmdline.present("-prune");
      b_weight = cmdline.present("-w");
      //enable gluino mode
      b_gluino = cmdline.present("-gluino");
      skipevt = cmdline.value<int>("-skipevt", 0); 
      maxevt = cmdline.value<int>("-n", -1);
      n_display = cmdline.value<int>("-display", 10);
      i_key = 0;
      /*
      if(!b_weight && input_filename.find("QCD")!=string::npos)
	{
	  cout<<"MESSAGE: QCD is found in the input,"
	    <<"enable event weighting?"<<endl;
	  char c_in=' ';
	  while(c_in !='y' && c_in!='n')
	    {
	      cout<<"(y or n): ";
	      cin>>c_in;
	    }
	  if(c_in=='y')
	    b_weight=true;
	  else
	    b_weight=false;
	}
      */

      if(b_gluino && b_weight)
	{
	  cout<<"WARNING: -gluino and -w cannot be enabled"
	    <<"at the same time\n"
	    <<"WARNING: -gluino diabled"<<endl;
	  b_gluino=false;
	}
      
      if(b_prune)
	cout<<"OPTIONS: Pruning is enabled."<<endl;

      if(b_trim)
	cout<<"OPTIONS: (default) Trimming is enabled."<<endl;

      if(b_pixel)
	cout<<"OPTIONS: (default) Pixelation is enabled."<<endl;
      
      if(b_weight)
	cout<<"OPTIONS: Event weighting is enabled."<<endl;
      
      if(b_gluino)
	cout<<"OPTIONS: Record gluino information."<<endl;

      //initialize basic histograms
      double c_cutweight=0;
      double c_cutweight_unweighted=0;

      //if input file is a root file, simple add it
      if(input_filename.find("root")!=string::npos)
	{
	  tree.Add(input_filename.c_str());
	  TFile c_file(input_filename.c_str());

	  //get the event weights that were skimmed away
	  if(skipevt==0){
	  TH1D* pc_nevt=(TH1D*) c_file.Get("nevt");
	  if(pc_nevt)
	    {
	      c_cutweight+=pc_nevt->GetBinContent(1)-
		pc_nevt->GetBinContent(2);
	    }

	  //get the event weights that were skimmed away
	  TH1D* pc_nevt_unweighted=(TH1D*) c_file.Get("nevt_unweighted");
	  if(pc_nevt_unweighted)
	    {
	      c_cutweight_unweighted+=pc_nevt_unweighted->GetBinContent(1)-
		pc_nevt_unweighted->GetBinContent(2);
	    }
	  c_file.Close();

	  }
	}
      //else it is a list of files, read it
      else
	{
	  ifstream input_filelist(input_filename.c_str());
	  if(!input_filelist.good())
	    {
	      cout<<"ERROR: input filelist cannot be read"<<endl;
	      throw;
	    }

	  string filename;
	  while (input_filelist.good())
	    {
	      getline(input_filelist, filename);
	      if(filename.find("root")!=string::npos)
		{
		  tree.Add(filename.c_str());
		  cout<<"File "<<filename<<" added"<<endl;
		  TFile c_file(input_filename.c_str());
		
		  if(skipevt==0){
		  TH1D* pc_nevt=(TH1D*) c_file.Get("nevt");
		  if(pc_nevt)
		    {
		      c_cutweight+=pc_nevt->GetBinContent(1)-
			pc_nevt->GetBinContent(2);
		    }

		  TH1D* pc_nevt_unweighted=(TH1D*) c_file.Get("nevt_unweighted");
		  if(pc_nevt_unweighted)
		    {
		      c_cutweight_unweighted+=pc_nevt_unweighted->GetBinContent(1)-
			pc_nevt_unweighted->GetBinContent(2);
		    }
		  
		  c_file.Close();
		  }//end if skipevent
		}
	    }
	  
	}
      cout<<"MESSAGE: "<<tree.GetEntries()<<" total entries"<<endl;
      if(tree.GetEntries()<= skipevt)
	{
	  cout<<"ERROR: too many events skipped"<<endl;
	  throw("NOT OUTPUT GENERATED");
	}

      else
	fileout=new TFile(output_filename.c_str(), "RECREATE"); 
	
      
      hist["nevt"]=new TH1D("nevt", "nevt", 10, 0, 10);
      hist["nevt_unweighted"]=new TH1D("nevt_unweighted", "nevt_unweighted",
				       10, 0, 10);

      hist["nevt"]->Fill(0.0, c_cutweight);
      hist["nevt_unweighted"]->Fill(0.0, c_cutweight_unweighted);
      

    }

  
  void tree_init()
  {

    

    //can take two kinds of trees
    //with weights, or without weights
    tree.SetMakeClass(1);
    
    //b_weight=true;
    if(!b_weight){
      weight = 1.0;
      
      tree.SetBranchAddress("Event", &Event_);
      tree.SetBranchAddress("Event.fUniqueID", Event_fUniqueID);
      tree.SetBranchAddress("Event.fBits", Event_fBits);
      tree.SetBranchAddress("Event.Number", Event_Number);
      tree.SetBranchAddress("Event_size", &Event_size);
      tree.SetBranchAddress("GenParticle", &GenParticle_);
      tree.SetBranchAddress("GenParticle.fUniqueID", fUniqueID);
      tree.SetBranchAddress("GenParticle.fBits", fBits);
      tree.SetBranchAddress("GenParticle.PID", PID);
      tree.SetBranchAddress("GenParticle.Status", Status);
      tree.SetBranchAddress("GenParticle.M1", M1);
      tree.SetBranchAddress("GenParticle.M2", M2);
      tree.SetBranchAddress("GenParticle.D1", D1);
      tree.SetBranchAddress("GenParticle.D2", D2);
      tree.SetBranchAddress("GenParticle.E", E);
      tree.SetBranchAddress("GenParticle.Px", Px);
      tree.SetBranchAddress("GenParticle.Py", Py);
      tree.SetBranchAddress("GenParticle.Pz", Pz);
      tree.SetBranchAddress("GenParticle.PT", Pt);
      tree.SetBranchAddress("GenParticle.Eta", Eta);
      tree.SetBranchAddress("GenParticle.Phi", Phi);
      tree.SetBranchAddress("GenParticle.Rapidity", Rapidity);
      tree.SetBranchAddress("GenParticle.T", T);
      tree.SetBranchAddress("GenParticle.X", X);
      tree.SetBranchAddress("GenParticle.Y", Y);
      tree.SetBranchAddress("GenParticle.Z", Z);
      tree.SetBranchAddress("GenParticle_size", &size);

      //test and see if the tree is the right format
      if(tree.GetLeaf("weight")!=NULL)
	  {
	    cout<<"MESSAGE: weight branch found, "
		<<"automatically enable weighting"<<endl;
	  }
	b_weight=true;
      }


    if(b_weight)
      {
	tree.SetBranchAddress("weight", &weight);
	tree.SetBranchAddress("size", &size);
	tree.SetBranchAddress("PID", PID);
	tree.SetBranchAddress("Status", Status);
	tree.SetBranchAddress("E", E);
	tree.SetBranchAddress("Px", Px);
	tree.SetBranchAddress("Py", Py);
	tree.SetBranchAddress("Pz", Pz);
	tree.SetBranchAddress("Pt", Pt);
	tree.SetBranchAddress("Eta", Eta);
	tree.SetBranchAddress("Phi", Phi);
	return;
      }




    

  }

  
  void reset()
  {
    InputJets.clear();
    jets.clear();
    leps.clear();
    met=PseudoJet();
    ttop_pt.clear();
    //reset all cell_E entries to 0
    //fastest way via memset
    memset(cell_E, 0, 
	   sizeof(cell_E[0][0])*cell_phi_size*cell_eta_size);
    if(clustSeq!=NULL)
      delete clustSeq;
  }

 protected:

  //options
  CmdLine cmdline;
  string input_filename, output_filename;
  string s_algorithm;
  int mass;
  double r;
  bool b_prune, b_weight, b_gluino, b_trim, b_pixel;
  int maxevt, skipevt;
  TFile* fileout;
  
  static const double stepsize_=.0997331;
  static const int cell_phi_size=63;
  static const int cell_eta_size=64;  
  static const double min_eta=-3.1914592;
  static const double max_eta=3.1914592;
    
  

  //Unit 4-vectors for each cell
  PseudoJet cell_vec[cell_eta_size][cell_phi_size];
  double cell_E[cell_eta_size][cell_phi_size];
  
  //number of diaply plots to plot
  //and counter for the display
  int n_display;
  int i_key;

  //tree variables
  const static int kMaxGenParticle=10000;
  

  //array of data;
  //DO NOT remove even unused variables, which may cause segfaults
  Double_t weight;

  //create sample display plot
  vector<string> key_E;
  vector<string> key_Pt;
  vector<string> key_gluino2;
  vector<string> key_gluino1;

  vector<vector<string> > key_jet;
  vector<int> key_entry;

  Int_t           Event_;
  UInt_t          Event_fUniqueID[1];
  UInt_t          Event_fBits[1];
  Long64_t        Event_Number[1];
  Int_t           Event_size;
  Int_t           GenParticle_;
  UInt_t          fUniqueID[kMaxGenParticle];
  UInt_t          fBits[kMaxGenParticle];
  Int_t           PID[kMaxGenParticle];
  Int_t           Status[kMaxGenParticle];
  Int_t           M1[kMaxGenParticle];
  Int_t           M2[kMaxGenParticle];
  Int_t           D1[kMaxGenParticle];
  Int_t           D2[kMaxGenParticle];
  Double_t        E[kMaxGenParticle];
  Double_t        Px[kMaxGenParticle];
  Double_t        Py[kMaxGenParticle];
  Double_t        Pz[kMaxGenParticle];
  Double_t        Pt[kMaxGenParticle];
  Double_t        Eta[kMaxGenParticle];
  Double_t        Phi[kMaxGenParticle];
  Double_t        Rapidity[kMaxGenParticle];
  Double_t        T[kMaxGenParticle];
  Double_t        X[kMaxGenParticle];
  Double_t        Y[kMaxGenParticle];
  Double_t        Z[kMaxGenParticle];
  Int_t           size;


  //private function to find daughters
  void find_daughters(set<int>& list)
  {
    bool continue_search=true;
    while(continue_search)
      {
	continue_search=false;
	set<int> listnew=list;
	set<int>::iterator it;
	
	for(it = list.begin(); it!=list.end(); ++it)
	  {
	    //cout<<"Daughters found !!: "<<D2[*it]<<" "<<D1[*it]<<endl;
	    if(D1[*it]>0)
	      {
		listnew.erase(*it);
		for(int i=D1[*it]; i<=D2[*it]; ++i)
		  {
		    if(M2[i]==0)
		      listnew.insert(i);
		  }
		continue_search=true;
	      }
	  }
	if(continue_search)
	  list=listnew;
      }
  }
  
};
#endif
